from asyncio import timeout

from docx import Document
from mylangchain.chains.retrieval_qa.base import RetrievalQA
from mylangchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.prompts import PromptTemplate
from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_milvus import Milvus
from langchain_core.runnables import RunnablePassthrough

# 定义变量
BASE_URL = "http://10.0.2.114:11434"
MILVUS_URL = "http://192.168.206.134:19530"
WORD_PATH = "./file/IoT Platform Portal前端框架说明文档.docx"
EMBEDDING_MODEL = "bge-m3:latest"

# # 读取word文件
# def loadFile(path):
#     doc = Document(path)
#     return "\n".join(p.text for p in doc.paragraphs if p.text.strip())
#
#
# text = loadFile(WORD_PATH)
#
# # 文本切块
# text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=20)
# docs = text_splitter.create_documents([text])
#
# # 向量化 & 存入Milvus
# embeddings = OllamaEmbeddings(base_url=BASE_URL, model=EMBEDDING_MODEL)
# vector_store = Milvus.from_documents(documents=docs, embedding=embeddings,
#                                      connection_args={"uri": MILVUS_URL},
#                                      collection_name="test_milvus",
#                                      drop_old=True)

embeddings = OllamaEmbeddings(base_url=BASE_URL, model=EMBEDDING_MODEL)
vector_store = Milvus(collection_name="test_milvus", connection_args={"uri": MILVUS_URL}, embedding_function=embeddings)

# 大模型检索
llm = ChatOllama(base_url=BASE_URL, model="deepseek-r1:7b", temperature=1.0)

prompt = PromptTemplate(
    input_variables=["context", "question"],
    template="""你是一个从业20年的技术架构师，你必须使用我提供的上下文作为资料去回答问题，如果上下文没有足够的信息, 回答"在提供的资料中未找到答案".

<context>
{context}
</context>

Question: {question}
"""
)

retriever = vector_store.as_retriever(search_kwargs={"k": 5})


def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


chain_milvus = ({"context": retriever | format_docs, "question": RunnablePassthrough()}
         | prompt
         | llm
         )

# for chunk in chain.stream("代码仓库地址是什么？"):
#     print(chunk.content, end="", flush=True)
